package cn.frank.framework.caesar.module.ai.openAi.rest;

import cn.frank.framework.caesar.module.ai.openAi.conf.arrangement.assistant.Assistant;
import dev.langchain4j.service.TokenStream;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.extensions.Extension;
import io.swagger.v3.oas.annotations.extensions.ExtensionProperty;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

/**
 * @description: TODO
 * @author: renjiangyang
 * @date 2025/6/7 14:43:51
 * @since 1.8
 */
@Tag(name = "2 会话-共享记忆", description = "session", extensions = {
        @Extension(properties = {@ExtensionProperty(name = "x-order", value = "2", parseValue = true)}
        )
})
@RestController
@RequestMapping("/ai/openAiMemory")
public class OpenAiChatMemory1RestController extends BaseOpenAiRestController {

    @Autowired
    Assistant assistant;

    @Operation(summary = "提问")
    @GetMapping(value = "/chat")
    public String chat(@RequestParam(value = "message", defaultValue = "你是谁") String message) {
        return assistant.chat(message);
    }

    @Operation(summary = "流式")
    @GetMapping(value = "/streamBak", produces = "text/event-stream;charset=UTF-8")
    public Flux<String> memoryStreamChat(@RequestParam(value = "message", defaultValue = "我是谁") String message, HttpServletResponse response) {
        TokenStream stream = assistant.stream(message);
        // return Flux.create(sink -> {
        //     stream.onPartialResponse(s -> sink.next(s))
        //             .onCompleteResponse(c -> sink.complete())
        //             .onError(sink::error)
        //             .start();
        //
        // });
        return Flux.create(sink -> {
            stream.onPartialResponse(s -> {
                        System.out.println("s = " + s);
                        sink.next("data: " + s + "\n\n");
                    })
                    .onCompleteResponse(c -> sink.complete())
                    .onError(sink::error)
                    .start();
        });
    }

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> memoryStreamChat(@RequestParam("message") String message) {
        TokenStream stream = assistant.stream(message);

        return Flux.create(sink -> {
            stream.onPartialResponse(s -> {
                        System.out.println("s = " + s);
                        sink.next(ServerSentEvent.builder(s).build());// 强制 flush
                    })
                    .onCompleteResponse(c -> {
                        sink.next(ServerSentEvent.builder("[DONE]").event("done").build());
                        sink.complete();
                    })
                    .onError(sink::error)
                    .start();
        });
    }


}
